{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e2eb2279-c8c3-4c9a-a0ba-154548bd4872",
   "metadata": {},
   "source": [
    "## 四种文档处理链\n",
    "- Stuff\n",
    "- Refine\n",
    "- Map reduce\n",
    "- Map re-mark"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8eeb5772-9a39-47fc-8b0d-25a642e78b9e",
   "metadata": {},
   "source": [
    "## 1. StuffChain\n",
    "最常见的文档链，将文档直接塞进prompt中，为LLM回答问题提供上下文资料，适合小文档场景"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71e61305-0f4e-4fd7-baa6-72aa929cd89d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains.combine_documents.stuff import StuffDocumentsChain\n",
    "from langchain.chains import LLMChain\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.document_loaders import  PyPDFLoader\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "\n",
    "prompt_template = \"\"\"对以下文字做简洁的总结:\n",
    "{text}\n",
    "\n",
    "简洁的总结:\"\"\"\n",
    "\n",
    "prompt = PromptTemplate.from_template(prompt_template)\n",
    "llm = ChatOpenAI(\n",
    "    temperature=0,\n",
    "    model=\"gpt-4-1106-preview\",\n",
    ")\n",
    "llm_chain = LLMChain(llm=llm, prompt=prompt)\n",
    "\n",
    "stuff_chain = StuffDocumentsChain(\n",
    "    llm_chain=llm_chain,\n",
    "    document_variable_name=\"text\",\n",
    ")\n",
    "\n",
    "loader = PyPDFLoader(\"loader.pdf\")\n",
    "#print(loader.load())\n",
    "docs = loader.load()\n",
    "print(stuff_chain.run(docs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78788c95-b725-49ec-91dd-0fe131c158af",
   "metadata": {},
   "outputs": [],
   "source": [
    "#使用预封装好的load_summarize_chain\n",
    "from langchain.document_loaders import  PyPDFLoader\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.chains.summarize import load_summarize_chain\n",
    "\n",
    "loader = PyPDFLoader(\"loader.pdf\")\n",
    "docs = loader.load()\n",
    "llm = ChatOpenAI(\n",
    "    temperature=0,\n",
    "    model=\"gpt-4-1106-preview\",\n",
    ")\n",
    "# 已经预封装定义好的chain\n",
    "chain = load_summarize_chain(\n",
    "    llm=llm,\n",
    "    chain_type=\"stuff\",\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "chain.run(docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12085c0c-a57a-4814-961d-94f5c448ec44",
   "metadata": {},
   "source": [
    "## 2. RefineChain\n",
    "通过循环输入文档并迭代更新答案来构建响应，一次只传递给LLM一个文档，将文档不断投喂，并产生各种中间答案，适合逻辑有上下文关联的文档，不适合交叉引用的文档"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dae516ab-94c9-4000-a150-f3cf5ac89476",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts import  PromptTemplate\n",
    "from langchain.document_loaders import PyPDFLoader\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "from langchain.chains.summarize import load_summarize_chain\n",
    "\n",
    "prompt_template = \"\"\"对以下文字做简洁的总结:\n",
    "{text}\n",
    "简洁的总结:\"\"\"\n",
    "prompt = PromptTemplate.from_template(prompt_template)\n",
    "\n",
    "refine_template = (\n",
    "    \"你的任务是产生最终摘要\\n\"\n",
    "    \"我们已经提供了一个到某个特定点的现有回答:{existing_answer}\\n\"\n",
    "    \"我们有机会通过下面的一些更多上下文来完善现有的回答(仅在需要时使用).\\n\"\n",
    "    \"------------\\n\"\n",
    "    \"{text}\\n\"\n",
    "    \"------------\\n\"\n",
    "    \"根据新的上下文，用中文完善原始回答.\\n\"\n",
    "    \"如果上下文没有用处,返回原始回答.\"\n",
    ")\n",
    "refine_prompt = PromptTemplate.from_template(refine_template)\n",
    "\n",
    "llm = ChatOpenAI(\n",
    "    temperature=0,\n",
    "    model=\"gpt-3.5-turbo\",\n",
    ")\n",
    "\n",
    "chain = load_summarize_chain(\n",
    "    llm=llm,\n",
    "    chain_type=\"refine\",\n",
    "    question_prompt=prompt,\n",
    "    refine_prompt = refine_prompt,\n",
    "    return_intermediate_steps=True,\n",
    "    input_key = \"documents\",\n",
    "    output_key = \"output_text\",\n",
    ")\n",
    "\n",
    "#load\n",
    "loader = PyPDFLoader(\"loader.pdf\")\n",
    "docs = loader.load()\n",
    "#split\n",
    "text_split = CharacterTextSplitter.from_tiktoken_encoder(\n",
    "    chunk_size = 1000,\n",
    "    chunk_overlap=0\n",
    ")\n",
    "split_docs = text_split.split_documents(docs)\n",
    "\n",
    "result = chain({\"documents\":split_docs},return_only_outputs=True)\n",
    "print(result[\"output_text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d4cb448-5613-4b5f-8d67-95006c4b5f5a",
   "metadata": {},
   "source": [
    "## Map reduce\n",
    "先将每个文档或文档块分别投喂给LLM，并得到结果集（Map步骤），然后通过一个文档合并链，获得一个输出结果（Reduce步骤）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98cfd91b-c9cb-457b-9983-7ffa9a44341f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import MapReduceDocumentsChain\n",
    "from langchain.chains import ReduceDocumentsChain\n",
    "from langchain.chains.combine_documents.stuff import StuffDocumentsChain\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.chains import LLMChain\n",
    "from langchain.document_loaders import PyPDFLoader\n",
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "\n",
    "#load pdf\n",
    "loader = PyPDFLoader(\"loader.pdf\")\n",
    "docs = loader.load()\n",
    "#split text\n",
    "text_splitter = CharacterTextSplitter.from_tiktoken_encoder(\n",
    "    chunk_size=1000,\n",
    "    chunk_overlap=0,\n",
    ")\n",
    "split_docs = text_splitter.split_documents(docs)\n",
    "#print(split_docs)\n",
    "\n",
    "#map chain\n",
    "map_template = \"\"\"对以下文字做简洁的总结:\n",
    "\"{content}\"\n",
    "简洁的总结:\"\"\"\n",
    "map_prompt = PromptTemplate.from_template(map_template)\n",
    "llm = ChatOpenAI(\n",
    "    temperature=0,\n",
    "    model=\\\"gpt-3.5-turbo\\\",\n",
    ")\n",
    "map_chain = LLMChain(\n",
    "    llm=llm,\n",
    "    prompt=map_prompt,\n",
    ")\n",
    "\n",
    "#reduce chain\n",
    "reduce_template = \"\"\"以下是一个摘要集合:\n",
    "{doc_summaries}\n",
    "将上述摘要与所有关键细节进行总结.\n",
    "总结:\"\"\"\n",
    "reduce_prompt = PromptTemplate.from_template(reduce_template)\n",
    "reduce_chain = LLMChain(\n",
    "    prompt=reduce_prompt,\n",
    "    llm=llm,\n",
    ")\n",
    "stuff_chain = StuffDocumentsChain(\n",
    "    llm_chain=reduce_chain,\n",
    "    document_variable_name=\"doc_summaries\",\n",
    ")\n",
    "reduce_final_chain = ReduceDocumentsChain(\n",
    "    combine_documents_chain=stuff_chain,\n",
    "    #超过4000个token就会切入到下一个stuff_chain\n",
    "    collapse_documents_chain=stuff_chain,\n",
    "    token_max=4000,\n",
    ")\n",
    "\n",
    "#map reduce chain\n",
    "map_reduce_chain = MapReduceDocumentsChain(\n",
    "    llm_chain=map_chain,\n",
    "    document_variable_name=\"content\",\n",
    "    reduce_documents_chain=reduce_final_chain,\n",
    ")\n",
    "\n",
    "\n",
    "summary = map_reduce_chain.run(split_docs)\n",
    "print(summary)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42e42b49-d14e-48cf-b70b-54c2c28a08ba",
   "metadata": {},
   "source": [
    "## Map re-rank\n",
    "先将每个文档或文档块投喂给LLM,并对每个文档或文档块生成问题的答案进行打分，然后将打分最高的文档或文档块作为最终答案返回"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26f44cc4-9abc-47be-9d38-2af9845eae0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import PyPDFLoader\n",
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.chains.qa_with_sources import load_qa_with_sources_chain\n",
    "from langchain.prompts import PromptTemplate\n",
    "\n",
    "llm = ChatOpenAI(temperature=0, model_name=\"gpt-3.5-turbo-16k\")\n",
    "#load\n",
    "loader = PyPDFLoader(\"loader.pdf\")\n",
    "docs = loader.load()\n",
    "#split\n",
    "text_splitter = CharacterTextSplitter.from_tiktoken_encoder(\n",
    "    chunk_size=500, chunk_overlap=0\n",
    ")\n",
    "split_docs = text_splitter.split_documents(docs)\n",
    "\n",
    "chain = load_qa_with_sources_chain(\n",
    "    ChatOpenAI(temperature=0), \n",
    "    chain_type=\"map_rerank\", \n",
    "    metadata_keys=['source'], \n",
    "    return_intermediate_steps=True\n",
    ")\n",
    "\n",
    "query = \"what is this document talk about?answer by chinese\"\n",
    "result = chain({\"input_documents\":split_docs,\"question\":query})\n",
    "result"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (langchain)",
   "language": "python",
   "name": "langchain-env"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
